import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.ensemble import IsolationForest
from matplotlib.ticker import MultipleLocator
from sklearn.linear_model import LinearRegression

def safe_min_g(g):
    if g.empty:
        return None
    return g.loc[g['Wspd'].idxmin()]

def baseline(a, c, x):
    return a * x ** 2 + c


ssp = "/home/duxiangyu/origindata/labeled2_data/"
sp = "/home/duxiangyu/origindata/labeled_data/"
for filename in os.listdir(sp):
    X_full = pd.read_csv(os.path.join(sp, filename))
    print(filename)
    
    X_normal = X_full[(X_full['anomaly_label'] == 1)]
    n_bins = 1000
    y_min = X_normal['Patv'].min()
    y_max = X_normal['Patv'].max()
    bins = np.linspace(y_min, y_max, n_bins + 1)
    X_normal['y_bins'] = pd.cut(X_normal['Patv'], bins=bins)

    grouped = X_normal.groupby('y_bins')
    max_x_per_bin = grouped.apply(safe_min_g)
    max_x_per_bin = max_x_per_bin.reset_index(drop=True)

    sample_set = max_x_per_bin.loc[max_x_per_bin['Wspd'].notnull(), [ 'Wspd', 'Patv']]
    sample_set['Wspd_q'] = sample_set['Wspd'] ** 2
    # print(sample_set)
    x_q = sample_set[['Wspd_q']].values
    y = sample_set['Patv'].values
    model_line = LinearRegression()
    model_line.fit(x_q, y)

    a = model_line.coef_[0]
    intercept = model_line.intercept_
    print(f'拟合公式：Patv ≈ {a:.4f}·Wspd² + {intercept:.4f}')
    X_full['residual'] = X_full['Patv'] - X_full.apply(lambda row: baseline(a, intercept, row['Wspd']), axis=1)
    X_full['anomaly_label'] = None

    # print(X_full)

    # 筛除第一类异常点，第三类异常点
    X_full.loc[X_full['Patv'] < 0, 'anomaly_label'] = -1 
    X_full.loc[(X_full['Patv'] < 1) & (X_full['Wspd'] > 2.5), 'anomaly_label'] = -1
    X_full.loc[(X_full['Patv'] < 1400) & (X_full['Wspd'] > 13), 'anomaly_label'] = -1
    # 标记好可能被误判的满发点
    X_full.loc[(X_full['Patv'] > 1400) & (X_full['Wspd'] > 8), 'anomaly_label'] = 1

    X = X_full.loc[X_full['anomaly_label'].isna(), ['residual']]
    # print(X)
    model = IsolationForest(n_estimators=256, contamination=0.03, max_samples=1000, random_state=42)
    model.fit(X)
    y_pred = model.predict(X)

    X_full.loc[X.index, 'anomaly_label'] = y_pred
    X_full.to_csv(os.path.join(ssp, filename), index=False)
    
    # 可视化
    plt.figure(figsize=(12, 9))
    plt.scatter(X_full.loc[:, 'Wspd'], X_full.loc[:, 'Patv'], c=X_full['anomaly_label'], cmap='coolwarm', edgecolor='k')
    plt.title("Isolation Forest")
    plt.xlabel("Wspd")
    plt.ylabel("Patv")
    plt.grid(True)

    
    # 设置刻度间隔（横坐标每 1 个单位、纵坐标每 100 个单位）
    plt.gca().xaxis.set_major_locator(MultipleLocator(1))
    plt.gca().yaxis.set_major_locator(MultipleLocator(100))

    plt.savefig(f'a2_{filename}.png')
    plt.close()

    # break